#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

// template<typename float, typename float>
class KernelCumsum_fp32 {

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;

    GlobalTensor<float> xGm;
    GlobalTensor<float> yGm;
    
    uint32_t minLength;
    uint32_t ALIGN_NUM;

    uint32_t* shape;
    uint32_t dim;

    uint32_t axesNum;

    uint32_t blockLength;
    uint32_t blockSize;
    uint32_t unitCount;
    uint32_t totalLength; //reduce数据个数
    uint32_t stepSize;
    uint32_t totalNum;  //所有数据个数
    
    uint32_t tileNum; //每组reduce数据计算的循环次数
    uint32_t tileLength; //每次tiling的数据长度
    uint32_t lasttileLength; //最后一次tiling的长度/ uint32_t remain_size;

    bool exclusive;
    bool reverse;

public:
    __aicore__  inline KernelCumsum_fp32() {}
    __aicore__  inline  void Init(GM_ADDR x, GM_ADDR axes, GM_ADDR y, uint32_t blockLength, uint32_t blockSize, uint32_t ALIGN_NUM, uint32_t* shape, uint32_t dim, bool exclusive, bool reverse)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");


        this->blockLength = blockLength;
        this->blockSize = blockSize;
        this->ALIGN_NUM = ALIGN_NUM;
        this->minLength = this->ALIGN_NUM;
        this->shape = shape;
        this->dim = dim;
        this->exclusive = exclusive;
        this->reverse = reverse;

        this->axesNum = (*(__gm__ int32_t *)axes);
        if(axesNum < 0)
        {
            axesNum = dim + axesNum;
        }
        
        this->totalLength = shape[axesNum];
        
        this->stepSize = 1;
        this->unitCount = 1;
        for(uint32_t i = 0; i < axesNum; i++)
        {
            this->unitCount *= shape[i];
        }

        for(uint32_t i = axesNum+1; i < dim; i++)
        {
            this->stepSize *= shape[i];
        }

        this->totalNum = unitCount * totalLength * stepSize;

        uint32_t totalLengthAligned = (totalLength + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

        this->tileNum = totalLengthAligned / this->blockSize;
        // 
        if(totalLengthAligned % blockSize == 0 || tileNum == 0) //reduce dim对齐之后，可以均分到各个tile，或不足一个tile
        {
            if(tileNum == 0)    //一次tiling的数据个数 大于 totalLengthAligned，也就是tile_num = 0时，tileLength=totalLengthAligned
            {
                this->tileNum = 1;
                this->tileLength = totalLengthAligned;
                this->lasttileLength = this->totalLength;
            }
            else //一次tiling的数据个数 小于 totalLengthAligned时，tileLength取最大值tiling_size * ALIGN_NUM
            {
                this->tileLength = this->blockSize;
                this->lasttileLength = this->tileLength;
            }
        }
        else //totalLengthAligned对齐后不能均分到每次tiling
        {
            this->tileNum = this->tileNum + 1;
            this->tileLength = this->blockSize; //tileLength 取最大值(ub)
            this->lasttileLength = totalLengthAligned % this->tileLength;
        }

        // printf("axesNum : %u\n", axesNum);
        // printf("totalLength : %u\n", totalLength);
        // printf("unitCount : %u\n", unitCount);
        // printf("stepSize : %u\n", stepSize);
        // printf("tileNum : %u\n", tileNum);
        // printf("tileLength : %u\n", tileLength);
        // printf("lasttileLength : %u\n", lasttileLength);
        // printf("totalNum : %u\n", totalNum);
        // printf("exclusive : %d\n", exclusive);
        // printf("reverse : %d\n", reverse);

        ASSERT(tileNum != 0 && "tile num can not be zero!");
    
        
        xGm.SetGlobalBuffer((__gm__ float *)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ float *)y, blockLength);

        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(float));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(float));
    }

    __aicore__ inline void Process()
    {
        if (stepSize != 1)
        {
            // 处理轴 != -1情况, 需要SetValue来设置xlocal
            ProcessAxes();
            return;
        }
        //最后一个维度 stepsize = 1
        uint32_t loopCount = this->tileNum;
        for (int32_t k = 0; k < this->unitCount; k++)
        {
            float lastsum = 0.0;
            for(int32_t i = 0; i < loopCount - 1; i++)
            {
                CopyIn(i,  k, tileLength);
                Compute(i, tileLength, lastsum);
                CopyOut(i, k, tileLength);
            }
            uint32_t lastlength = reverse == false ? lasttileLength : totalLength - (tileNum - 1) * tileLength;
            CopyIn(tileNum - 1,  k, lastlength);
            Compute(tileNum - 1, totalLength - (tileNum - 1) * tileLength, lastsum);
            CopyOut(tileNum - 1, k, lastlength);
        }
    }


    __aicore__ inline void CopyIn(int32_t progress, int k, uint32_t length)
    {
        LocalTensor<float> xLocal = inQueueX.AllocTensor<float>();
        if(reverse == false)
        {
            uint32_t startIndex = k  * this->totalLength + progress * tileLength;
            DataCopy(xLocal, xGm[startIndex], length);
            inQueueX.EnQue(xLocal);
        }
        else if(reverse == true)
        {
            uint32_t startIndex = k  * this->totalLength + (this->totalLength -  (progress + 1) * tileLength);
            // startIndex = startIndex > 0 ? startIndex : 0;

            if(progress == this->tileNum - 1)   //最后一次tiling 且无法datacopy, 先copy再覆盖 ，因为lasttileLength 一定小于tileLength
            {
                startIndex = k  * this->totalLength;    //倒序最后一次tiling的start为 “0”  : *,*,*,*,0,0,0
                DataCopy(xLocal, xGm[startIndex], lasttileLength);
                for (int i = length; i < this->lasttileLength; i++)      //0 padding
                {
                    xLocal.SetValue(i, (float)0);
                }
            }
            else
            {
                DataCopy(xLocal, xGm[startIndex], length);
            }

            inQueueX.EnQue(xLocal);
        }
    }


    __aicore__ inline void Compute(int32_t progress, uint32_t length, float& lastsum)
    {   
        if(reverse == false)
        {
            LocalTensor<float> xLocal = inQueueX.DeQue<float>();
            LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();

            Adds(yLocal, xLocal, (float)0.0f, length);
            
            if(exclusive)
            {
                if(length > 1)
                {
                    yLocal.SetValue(0, (float)0.0f);
                    yLocal.SetValue(1, xLocal.GetValue(0));
                    for(int i = 2; i < length; i++)
                    {
                        yLocal.SetValue(i, yLocal.GetValue(i-1) + xLocal.GetValue(i-1));
                    }
                }
                else    // length = 1
                {
                    yLocal.SetValue(0, (float)0);
                }
            }
            else
            {   
                for(int i = 1; i < length; i++)
                {
                    yLocal.SetValue(i, yLocal.GetValue(i-1) + xLocal.GetValue(i));
                }
            }
            Adds(yLocal, yLocal, lastsum, length);
            lastsum = yLocal.GetValue(length - 1);

            outQueueY.EnQue<float>(yLocal);
            inQueueX.FreeTensor(xLocal);
        }
        else if (reverse == true)
        {
            LocalTensor<float> xLocal = inQueueX.DeQue<float>();
            LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();

            Adds(yLocal, xLocal, (float)0.0f, length);
            
            if(exclusive)
            {
                if(length > 1)
                    {
                        yLocal.SetValue(length - 1, (float)0.0f);
                        yLocal.SetValue(length - 2, xLocal.GetValue(length - 1));
                        for(int i = length - 3; i > -1; i--)
                        {
                            yLocal.SetValue(i, xLocal.GetValue(i+1) + yLocal.GetValue(i+1));
                        }
                    }
                    else    // length = 1
                    {
                        yLocal.SetValue(length - 1, (float)0.0f);
                    }
            }
            else
            {   
                for(int i = length - 2; i > -1; i--)
                {
                    yLocal.SetValue(i, yLocal.GetValue(i+1) + xLocal.GetValue(i));
                }
            }
            Adds(yLocal, yLocal, lastsum, length);
            lastsum = yLocal.GetValue(0);

            outQueueY.EnQue<float>(yLocal);
            inQueueX.FreeTensor(xLocal);
        }   
    }


    __aicore__ inline void CopyOut(int32_t progress, int k, uint32_t length)
    {
        if(reverse ==false)
        {
            uint32_t startIndex = k * this->totalLength + progress * tileLength;
            LocalTensor<float> yLocal = outQueueY.DeQue<float>();
            DataCopy(yGm[startIndex], yLocal, length);    
            outQueueY.FreeTensor(yLocal);    
        }
        else if(reverse == true)
        {
            LocalTensor<float> yLocal = outQueueY.DeQue<float>();
            if(progress == this->tileNum - 1)   //最后一次tiling
            {
                
                uint32_t startIndex = k  * this->totalLength;
                if(lasttileLength != length)    //存在0padding的数据，只能setvalue拷贝
                    for(int i = 0; i < length; i++)
                    {
                        yGm.SetValue(startIndex + i, yLocal.GetValue(i));
                    }
                else    //不存在0拷贝的数据
                    DataCopy(yGm[startIndex], yLocal, length);
            }
            else
            {
                uint32_t startIndex = k  * this->totalLength + (this->totalLength -  (progress + 1) * tileLength);
                LocalTensor<float> yLocal = outQueueY.DeQue<float>();
                DataCopy(yGm[startIndex], yLocal, length);
            }
            outQueueY.FreeTensor(yLocal);
        }        
    }


    __aicore__ inline void ProcessAxes()
    {
        uint32_t loopCount = this->tileNum;
        for (int32_t k = 0; k < this->unitCount; k++)
        {
            for (int32_t j = 0; j < this->stepSize; j++)
            {
                float lastsum = 0.0;
                for(int32_t i = 0; i < loopCount - 1; i++)
                {
                    CopyInAxes(i, j, k, tileLength);
                    Compute(i, tileLength, lastsum);
                    CopyOutAxes(i, j, k, tileLength);
                }
                CopyInAxes(tileNum - 1, j, k, totalLength - (tileNum -1)*tileLength);
                Compute(tileNum -1, totalLength - (tileNum -1)*tileLength, lastsum);
                CopyOutAxes(tileNum - 1, j, k, totalLength - (tileNum -1)*tileLength);
            }   
        }
    }


    __aicore__ inline void CopyInAxes(int32_t progress, int j, int k, uint32_t length)
    {
        if(reverse == false)
        {
            LocalTensor<float> xLocal = inQueueX.AllocTensor<float>();
            uint32_t startIndex = k * this->stepSize * this->totalLength +  progress * tileLength * this->stepSize + j;
            for (int i = 0; i < length; i++)
            {
                xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
            }
            inQueueX.EnQue(xLocal);
        }
        else if(reverse == true)
        {
            LocalTensor<float> xLocal = inQueueX.AllocTensor<float>();
            uint32_t startIndex = k * this->stepSize * this->totalLength + (this->totalLength - (progress + 1) * tileLength) * this->stepSize + j;

            if(progress == this->tileNum - 1)   //最后一次tiling
            {   
                startIndex = k * this->stepSize * this->totalLength + j;    // *,*,*,*,0,0
                for (int i = 0; i < length; i++)
                {
                    xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
                }
                for (int i = length; i < this->lasttileLength; i++)     //0 padding
                {
                    xLocal.SetValue(i, (float)0);
                }
            }
            else
            {
                for (int i = 0; i < length; i++)
                {
                    xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
                }
            }
            inQueueX.EnQue(xLocal);
        }
    }


    __aicore__ inline void CopyOutAxes(int32_t progress, int j, int k, uint32_t length)
    {
        if(reverse == false)
        {
            uint32_t startIndex = k * this->stepSize * this->totalLength + progress * tileLength * this->stepSize + j;
            LocalTensor<float> yLocal = outQueueY.DeQue<float>();
            for (int i = 0; i < length; i++)
            {
                yGm.SetValue(startIndex + i * stepSize, yLocal.GetValue(i));
            }
            outQueueY.FreeTensor(yLocal);
        }
        else if (reverse == true)
        {
            LocalTensor<float> yLocal = outQueueY.DeQue<float>();
            uint32_t startIndex;
            if(progress == this->tileNum - 1)   //最后一次tiling
            {
                startIndex = k * this->stepSize * this->totalLength + j;
            }
            else
            {
                startIndex = k * this->stepSize * this->totalLength + (this->totalLength - (progress + 1) * tileLength) * this->stepSize + j;
            }
            for(int i = 0; i < length; i++)
            {
                yGm.SetValue(startIndex + i * stepSize, yLocal.GetValue(i));
            }
            outQueueY.FreeTensor(yLocal);
        }    
    }
};


class KernelCumsum_fp16 {

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<TPosition::VECCALC> xfBuf, yfBuf;

    GlobalTensor<half> xGm;
    GlobalTensor<half> yGm;
    
    uint32_t minLength;
    uint32_t ALIGN_NUM;

    uint32_t* shape;
    uint32_t dim;

    uint32_t axesNum;

    uint32_t blockLength;
    uint32_t blockSize;
    uint32_t unitCount;
    uint32_t totalLength; //reduce数据个数
    uint32_t stepSize;
    uint32_t totalNum;  //所有数据个数
    
    uint32_t tileNum; //每组reduce数据计算的循环次数
    uint32_t tileLength; //每次tiling的数据长度
    uint32_t lasttileLength; //最后一次tiling的长度/ uint32_t remain_size;

    bool exclusive;
    bool reverse;

public:
    __aicore__  inline KernelCumsum_fp16() {}
    __aicore__  inline  void Init(GM_ADDR x, GM_ADDR axes, GM_ADDR y, uint32_t blockLength, uint32_t blockSize, uint32_t ALIGN_NUM, uint32_t* shape, uint32_t dim, bool exclusive, bool reverse)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");


        this->blockLength = blockLength;
        this->blockSize = blockSize;
        this->ALIGN_NUM = ALIGN_NUM;
        this->minLength = this->ALIGN_NUM;
        this->shape = shape;
        this->dim = dim;
        this->exclusive = exclusive;
        this->reverse = reverse;

        this->axesNum = (*(__gm__ int32_t *)axes);
        if(axesNum < 0)
        {
            axesNum = dim + axesNum;
        }
        
        this->totalLength = shape[axesNum];
        
        this->stepSize = 1;
        this->unitCount = 1;
        for(uint32_t i = 0; i < axesNum; i++)
        {
            this->unitCount *= shape[i];
        }

        for(uint32_t i = axesNum+1; i < dim; i++)
        {
            this->stepSize *= shape[i];
        }

        this->totalNum = unitCount * totalLength * stepSize;

        uint32_t totalLengthAligned = (totalLength + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

        this->tileNum = totalLengthAligned / this->blockSize;
        // 
        if(totalLengthAligned % blockSize == 0 || tileNum == 0) //reduce dim对齐之后，可以均分到各个tile，或不足一个tile
        {
            if(tileNum == 0)    //一次tiling的数据个数 大于 totalLengthAligned，也就是tile_num = 0时，tileLength=totalLengthAligned
            {
                this->tileNum = 1;
                this->tileLength = totalLengthAligned;
                this->lasttileLength = this->totalLength;
            }
            else //一次tiling的数据个数 小于 totalLengthAligned时，tileLength取最大值tiling_size * ALIGN_NUM
            {
                this->tileLength = this->blockSize;
                this->lasttileLength = this->tileLength;
            }
        }
        else //totalLengthAligned对齐后不能均分到每次tiling
        {
            this->tileNum = this->tileNum + 1;
            this->tileLength = this->blockSize; //tileLength 取最大值(ub)
            this->lasttileLength = totalLengthAligned % this->tileLength;
        }

        // printf("axesNum : %u\n", axesNum);
        // printf("totalLength : %u\n", totalLength);
        // printf("unitCount : %u\n", unitCount);
        // printf("stepSize : %u\n", stepSize);
        // printf("tileNum : %u\n", tileNum);
        // printf("tileLength : %u\n", tileLength);
        // printf("lasttileLength : %u\n", lasttileLength);
        // printf("totalNum : %u\n", totalNum);
        // printf("exclusive : %d\n", exclusive);
        // printf("reverse : %d\n", reverse);

        ASSERT(tileNum != 0 && "tile num can not be zero!");
    
        
        xGm.SetGlobalBuffer((__gm__ half *)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ half *)y, blockLength);

        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(half));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(half));
        pipe.InitBuffer(xfBuf, this->tileLength * sizeof(float));
        pipe.InitBuffer(yfBuf, this->tileLength * sizeof(float));
    }

    __aicore__ inline void Process()
    {
        if (stepSize != 1)
        {
            // 处理轴 != -1情况, 需要SetValue来设置xlocal
            ProcessAxes();
            return;
        }
        //最后一个维度 stepsize = 1
        uint32_t loopCount = this->tileNum;
        for (int32_t k = 0; k < this->unitCount; k++)
        {
            half lastsum = 0.0;
            for(int32_t i = 0; i < loopCount - 1; i++)
            {
                CopyIn(i,  k, tileLength);
                Compute(i, tileLength, lastsum);
                CopyOut(i, k, tileLength);
            }
            uint32_t lastlength = reverse == false ? lasttileLength : totalLength - (tileNum - 1) * tileLength;
            CopyIn(tileNum - 1,  k, lastlength);
            Compute(tileNum - 1, totalLength - (tileNum - 1) * tileLength, lastsum);
            CopyOut(tileNum - 1, k, lastlength);
        }
    }

    __aicore__ inline void CopyIn(int32_t progress, int k, uint32_t length)
    {
        LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
        if(reverse == false)
        {
            uint32_t startIndex = k  * this->totalLength + progress * tileLength;
            DataCopy(xLocal, xGm[startIndex], length);
            inQueueX.EnQue(xLocal);
        }
        else if(reverse == true)
        {
            uint32_t startIndex = k  * this->totalLength + (this->totalLength -  (progress + 1) * tileLength);
            // startIndex = startIndex > 0 ? startIndex : 0;

            if(progress == this->tileNum - 1)   //最后一次tiling 且无法datacopy, 先copy再覆盖 ，因为lasttileLength 一定小于tileLength
            {
                startIndex = k  * this->totalLength;    //倒序最后一次tiling的start为 “0”  : *,*,*,*,0,0,0
                DataCopy(xLocal, xGm[startIndex], lasttileLength);
                for (int i = length; i < this->lasttileLength; i++)      //0 padding
                {
                    xLocal.SetValue(i, (half)0);
                }
            }
            else
            {
                DataCopy(xLocal, xGm[startIndex], length);
            }

            inQueueX.EnQue(xLocal);
        }
    }


    __aicore__ inline void Compute(int32_t progress, uint32_t length, half& lastsum)
    {   
        if(reverse == false)
        {
            LocalTensor<half> xLocal_f = inQueueX.DeQue<half>();
            LocalTensor<half> yLocal_f = outQueueY.AllocTensor<half>();

            // LocalTensor<float> yLocal_f = yfBuf.Get<float>();
            // LocalTensor<float> xLocal_f = xfBuf.Get<float>();

            // Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, length);

            Adds(yLocal_f, xLocal_f, (half)0.0f, length);
            
            if(exclusive)
            {
                if(length > 1)
                {
                    yLocal_f.SetValue(0, (half)0.0f);
                    yLocal_f.SetValue(1, xLocal_f.GetValue(0));
                    for(int i = 2; i < length; i++)
                    {
                        yLocal_f.SetValue(i, (half)((float)yLocal_f.GetValue(i-1) + (float)xLocal_f.GetValue(i-1)));
                    }
                }
                else    // length = 1
                {
                    yLocal_f.SetValue(0, (half)0.0);
                }
            }
            else
            {   
                for(int i = 1; i < length; i++)
                {
                    yLocal_f.SetValue(i, (half)((float)yLocal_f.GetValue(i-1) + (float)xLocal_f.GetValue(i)));
                }
            }
            Adds(yLocal_f, yLocal_f, lastsum, length);
            lastsum = yLocal_f.GetValue(length - 1);

            // Cast(yLocal, yLocal_f, RoundMode::CAST_NONE, length);

            outQueueY.EnQue<half>(yLocal_f);
            inQueueX.FreeTensor(xLocal_f);
        }
        else if (reverse == true)
        {
            LocalTensor<half> xLocal_f = inQueueX.DeQue<half>();
            LocalTensor<half> yLocal_f = outQueueY.AllocTensor<half>();

            // LocalTensor<float> yLocal_f = yfBuf.Get<float>();
            // LocalTensor<float> xLocal_f = xfBuf.Get<float>();

            // Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, length);

            Adds(yLocal_f, xLocal_f, (half)0.0f, length);
            
            if(exclusive)
            {
                if(length > 1)
                    {
                        yLocal_f.SetValue(length - 1, (half)0.0);
                        yLocal_f.SetValue(length - 2, xLocal_f.GetValue(length - 1));
                        for(int i = length - 3; i > -1; i--)
                        {
                            yLocal_f.SetValue(i, (half)((float)xLocal_f.GetValue(i+1) + (float)yLocal_f.GetValue(i+1)));
                        }
                    }
                    else    // length = 1
                    {
                        yLocal_f.SetValue(length - 1, (half)0.0);
                    }
            }
            else
            {   
                for(int i = length - 2; i > -1; i--)
                {
                    yLocal_f.SetValue(i, (half)((float)yLocal_f.GetValue(i+1) + (float)xLocal_f.GetValue(i)));
                }
            }
            Adds(yLocal_f, yLocal_f, lastsum, length);
            lastsum = yLocal_f.GetValue(0);

            // Cast(yLocal, yLocal_f, RoundMode::CAST_NONE, length);
            outQueueY.EnQue<half>(yLocal_f);
            inQueueX.FreeTensor(xLocal_f);
        }

    }


    __aicore__ inline void CopyOut(int32_t progress, int k, uint32_t length)
    {
        if(reverse ==false)
        {
            uint32_t startIndex = k * this->totalLength + progress * tileLength;
            LocalTensor<half> yLocal = outQueueY.DeQue<half>();
            DataCopy(yGm[startIndex], yLocal, length);    
            outQueueY.FreeTensor(yLocal);    
        }
        else if(reverse == true)
        {
            LocalTensor<half> yLocal = outQueueY.DeQue<half>();
            if(progress == this->tileNum - 1)   //最后一次tiling
            {
                
                uint32_t startIndex = k  * this->totalLength;
                if(lasttileLength != length)    //存在0padding的数据，只能setvalue拷贝
                    for(int i = 0; i < length; i++)
                    {
                        yGm.SetValue(startIndex + i, yLocal.GetValue(i));
                    }
                else    //不存在0拷贝的数据
                    DataCopy(yGm[startIndex], yLocal, length);
            }
            else
            {
                uint32_t startIndex = k  * this->totalLength + (this->totalLength -  (progress + 1) * tileLength);
                LocalTensor<half> yLocal = outQueueY.DeQue<half>();
                DataCopy(yGm[startIndex], yLocal, length);
            }
            outQueueY.FreeTensor(yLocal);
        }        
    }


    __aicore__ inline void ProcessAxes()
    {
        uint32_t loopCount = this->tileNum;
        for (int32_t k = 0; k < this->unitCount; k++)
        {
            for (int32_t j = 0; j < this->stepSize; j++)
            {
                half lastsum = 0.0;
                for(int32_t i = 0; i < loopCount - 1; i++)
                {
                    CopyInAxes(i, j, k, tileLength);
                    Compute(i, tileLength, lastsum);
                    CopyOutAxes(i, j, k, tileLength);
                }
                CopyInAxes(tileNum - 1, j, k, totalLength - (tileNum -1)*tileLength);
                Compute(tileNum -1, totalLength - (tileNum -1)*tileLength, lastsum);
                CopyOutAxes(tileNum - 1, j, k, totalLength - (tileNum -1)*tileLength);
            }   
        }
    }


    __aicore__ inline void CopyInAxes(int32_t progress, int j, int k, uint32_t length)
    {
        if(reverse == false)
        {
            LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
            uint32_t startIndex = k * this->stepSize * this->totalLength +  progress * tileLength * this->stepSize + j;
            for (int i = 0; i < length; i++)
            {
                xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
            }
            inQueueX.EnQue(xLocal);
        }
        else if(reverse == true)
        {
            LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
            uint32_t startIndex = k * this->stepSize * this->totalLength + (this->totalLength - (progress + 1) * tileLength) * this->stepSize + j;

            if(progress == this->tileNum - 1)   //最后一次tiling
            {   
                startIndex = k * this->stepSize * this->totalLength + j;    // *,*,*,*,0,0
                for (int i = 0; i < length; i++)
                {
                    xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
                }
                for (int i = length; i < this->lasttileLength; i++)     //0 padding
                {
                    xLocal.SetValue(i, (half)0);
                }
            }
            else
            {
                for (int i = 0; i < length; i++)
                {
                    xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
                }
            }
            inQueueX.EnQue(xLocal);
        }        
    }

    __aicore__ inline void CopyOutAxes(int32_t progress, int j, int k, uint32_t length)
    {
        if(reverse == false)
        {
            uint32_t startIndex = k * this->stepSize * this->totalLength + progress * tileLength * this->stepSize + j;
            LocalTensor<half> yLocal = outQueueY.DeQue<half>();
            for (int i = 0; i < length; i++)
            {
                yGm.SetValue(startIndex + i * stepSize, yLocal.GetValue(i));
            }
            outQueueY.FreeTensor(yLocal);
        }
        else if (reverse == true)
        {
            LocalTensor<half> yLocal = outQueueY.DeQue<half>();
            uint32_t startIndex;
            if(progress == this->tileNum - 1)   //最后一次tiling
            {
                startIndex = k * this->stepSize * this->totalLength + j;
            }
            else
            {
                startIndex = k * this->stepSize * this->totalLength + (this->totalLength - (progress + 1) * tileLength) * this->stepSize + j;
            }
            for(int i = 0; i < length; i++)
            {
                yGm.SetValue(startIndex + i * stepSize, yLocal.GetValue(i));
            }
            outQueueY.FreeTensor(yLocal);
        }    
    }
};


class KernelCumsum_int32 {

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;

    GlobalTensor<int32_t> xGm;
    GlobalTensor<int32_t> yGm;
    
    uint32_t minLength;
    uint32_t ALIGN_NUM;

    uint32_t* shape;
    uint32_t dim;

    uint32_t axesNum;

    uint32_t blockLength;
    uint32_t blockSize;
    uint32_t unitCount;
    uint32_t totalLength; //reduce数据个数
    uint32_t stepSize;
    uint32_t totalNum;  //所有数据个数
    
    uint32_t tileNum; //每组reduce数据计算的循环次数
    uint32_t tileLength; //每次tiling的数据长度
    uint32_t lasttileLength; //最后一次tiling的长度/ uint32_t remain_size;

    bool exclusive;
    bool reverse;

public:
    __aicore__  inline KernelCumsum_int32() {}
    __aicore__  inline  void Init(GM_ADDR x, GM_ADDR axes, GM_ADDR y, uint32_t blockLength, uint32_t blockSize, uint32_t ALIGN_NUM, uint32_t* shape, uint32_t dim, bool exclusive, bool reverse)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");


        this->blockLength = blockLength;
        this->blockSize = blockSize;
        this->ALIGN_NUM = ALIGN_NUM;
        this->minLength = this->ALIGN_NUM;
        this->shape = shape;
        this->dim = dim;
        this->exclusive = exclusive;
        this->reverse = reverse;

        this->axesNum = (*(__gm__ int32_t *)axes);
        if(axesNum < 0)
        {
            axesNum = dim + axesNum;
        }
        
        this->totalLength = shape[axesNum];
        
        this->stepSize = 1;
        this->unitCount = 1;
        for(uint32_t i = 0; i < axesNum; i++)
        {
            this->unitCount *= shape[i];
        }

        for(uint32_t i = axesNum+1; i < dim; i++)
        {
            this->stepSize *= shape[i];
        }

        this->totalNum = unitCount * totalLength * stepSize;

        uint32_t totalLengthAligned = (totalLength + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

        this->tileNum = totalLengthAligned / this->blockSize;
        // 
        if(totalLengthAligned % blockSize == 0 || tileNum == 0) //reduce dim对齐之后，可以均分到各个tile，或不足一个tile
        {
            if(tileNum == 0)    //一次tiling的数据个数 大于 totalLengthAligned，也就是tile_num = 0时，tileLength=totalLengthAligned
            {
                this->tileNum = 1;
                this->tileLength = totalLengthAligned;
                this->lasttileLength = this->totalLength;
            }
            else //一次tiling的数据个数 小于 totalLengthAligned时，tileLength取最大值tiling_size * ALIGN_NUM
            {
                this->tileLength = this->blockSize;
                this->lasttileLength = this->tileLength;
            }
        }
        else //totalLengthAligned对齐后不能均分到每次tiling
        {
            this->tileNum = this->tileNum + 1;
            this->tileLength = this->blockSize; //tileLength 取最大值(ub)
            this->lasttileLength = totalLengthAligned % this->tileLength;
        }

        // printf("axesNum : %u\n", axesNum);
        // printf("totalLength : %u\n", totalLength);
        // printf("unitCount : %u\n", unitCount);
        // printf("stepSize : %u\n", stepSize);
        // printf("tileNum : %u\n", tileNum);
        // printf("tileLength : %u\n", tileLength);
        // printf("lasttileLength : %u\n", lasttileLength);
        // printf("totalNum : %u\n", totalNum);
        // printf("exclusive : %d\n", exclusive);
        // printf("reverse : %d\n", reverse);

        ASSERT(tileNum != 0 && "tile num can not be zero!");
    
        
        xGm.SetGlobalBuffer((__gm__ int32_t *)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ int32_t *)y, blockLength);

        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(int32_t));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(int32_t));
    }

    __aicore__ inline void Process()
    {
        if (stepSize != 1)
        {
            // 处理轴 != -1情况, 需要SetValue来设置xlocal
            ProcessAxes();
            return;
        }
        //最后一个维度 stepsize = 1
        uint32_t loopCount = this->tileNum;
        for (int32_t k = 0; k < this->unitCount; k++)
        {
            int32_t lastsum = 0.0;
            for(int32_t i = 0; i < loopCount - 1; i++)
            {
                CopyIn(i,  k, tileLength);
                Compute(i, tileLength, lastsum);
                CopyOut(i, k, tileLength);
            }
            uint32_t lastlength = reverse == false ? lasttileLength : totalLength - (tileNum - 1) * tileLength;
            CopyIn(tileNum - 1,  k, lastlength);
            Compute(tileNum - 1, totalLength - (tileNum - 1) * tileLength, lastsum);
            CopyOut(tileNum - 1, k, lastlength);
        }
    }


    __aicore__ inline void CopyIn(int32_t progress, int k, uint32_t length)
    {
        LocalTensor<int32_t> xLocal = inQueueX.AllocTensor<int32_t>();
        if(reverse == false)
        {
            uint32_t startIndex = k  * this->totalLength + progress * tileLength;
            DataCopy(xLocal, xGm[startIndex], length);
            inQueueX.EnQue(xLocal);
        }
        else if(reverse == true)
        {
            uint32_t startIndex = k  * this->totalLength + (this->totalLength -  (progress + 1) * tileLength);
            // startIndex = startIndex > 0 ? startIndex : 0;

            if(progress == this->tileNum - 1)   //最后一次tiling 且无法datacopy, 先copy再覆盖 ，因为lasttileLength 一定小于tileLength
            {
                startIndex = k  * this->totalLength;    //倒序最后一次tiling的start为 “0”  : *,*,*,*,0,0,0
                DataCopy(xLocal, xGm[startIndex], lasttileLength);
                for (int i = length; i < this->lasttileLength; i++)      //0 padding
                {
                    xLocal.SetValue(i, (int32_t)0);
                }
            }
            else
            {
                DataCopy(xLocal, xGm[startIndex], length);
            }

            inQueueX.EnQue(xLocal);
        }
    }


    __aicore__ inline void Compute(int32_t progress, uint32_t length, int32_t& lastsum)
    {   
        if(reverse == false)
        {
            LocalTensor<int32_t> xLocal = inQueueX.DeQue<int32_t>();
            LocalTensor<int32_t> yLocal = outQueueY.AllocTensor<int32_t>();

            Adds(yLocal, xLocal, (int32_t)0, length);
            
            if(exclusive)
            {
                if(length > 1)
                {
                    yLocal.SetValue(0, (int32_t)0);
                    yLocal.SetValue(1, xLocal.GetValue(0));
                    for(int i = 2; i < length; i++)
                    {
                        yLocal.SetValue(i, yLocal.GetValue(i-1) + xLocal.GetValue(i-1));
                    }
                }
                else    // length = 1
                {
                    yLocal.SetValue(0, (int32_t)0);
                }
            }
            else
            {   
                for(int i = 1; i < length; i++)
                {
                    yLocal.SetValue(i, yLocal.GetValue(i-1) + xLocal.GetValue(i));
                }
            }
            Adds(yLocal, yLocal, lastsum, length);
            lastsum = yLocal.GetValue(length - 1);

            outQueueY.EnQue<int32_t>(yLocal);
            inQueueX.FreeTensor(xLocal);
        }
        else if (reverse == true)
        {
            LocalTensor<int32_t> xLocal = inQueueX.DeQue<int32_t>();
            LocalTensor<int32_t> yLocal = outQueueY.AllocTensor<int32_t>();

            Adds(yLocal, xLocal, (int32_t)0, length);
            
            if(exclusive)
            {
                if(length > 1)
                    {
                        yLocal.SetValue(length - 1, (int32_t)0);
                        yLocal.SetValue(length - 2, xLocal.GetValue(length - 1));
                        for(int i = length - 3; i > -1; i--)
                        {
                            yLocal.SetValue(i, xLocal.GetValue(i+1) + yLocal.GetValue(i+1));
                        }
                    }
                    else    // length = 1
                    {
                        yLocal.SetValue(length - 1, (int32_t)0);
                    }
            }
            else
            {   
                for(int i = length - 2; i > -1; i--)
                {
                    yLocal.SetValue(i, yLocal.GetValue(i+1) + xLocal.GetValue(i));
                }
            }
            Adds(yLocal, yLocal, lastsum, length);
            lastsum = yLocal.GetValue(0);

            outQueueY.EnQue<int32_t>(yLocal);
            inQueueX.FreeTensor(xLocal);
        }   
    }


    __aicore__ inline void CopyOut(int32_t progress, int k, uint32_t length)
    {
        if(reverse ==false)
        {
            uint32_t startIndex = k * this->totalLength + progress * tileLength;
            LocalTensor<int32_t> yLocal = outQueueY.DeQue<int32_t>();
            DataCopy(yGm[startIndex], yLocal, length);    
            outQueueY.FreeTensor(yLocal);    
        }
        else if(reverse == true)
        {
            LocalTensor<int32_t> yLocal = outQueueY.DeQue<int32_t>();
            if(progress == this->tileNum - 1)   //最后一次tiling
            {
                
                uint32_t startIndex = k  * this->totalLength;
                if(lasttileLength != length)    //存在0padding的数据，只能setvalue拷贝
                    for(int i = 0; i < length; i++)
                    {
                        yGm.SetValue(startIndex + i, yLocal.GetValue(i));
                    }
                else    //不存在0拷贝的数据
                    DataCopy(yGm[startIndex], yLocal, length);
            }
            else
            {
                uint32_t startIndex = k  * this->totalLength + (this->totalLength -  (progress + 1) * tileLength);
                LocalTensor<int32_t> yLocal = outQueueY.DeQue<int32_t>();
                DataCopy(yGm[startIndex], yLocal, length);
            }
            outQueueY.FreeTensor(yLocal);
        }        
    }


    __aicore__ inline void ProcessAxes()
    {
        uint32_t loopCount = this->tileNum;
        for (int32_t k = 0; k < this->unitCount; k++)
        {
            for (int32_t j = 0; j < this->stepSize; j++)
            {
                int32_t lastsum = 0.0;
                for(int32_t i = 0; i < loopCount - 1; i++)
                {
                    CopyInAxes(i, j, k, tileLength);
                    Compute(i, tileLength, lastsum);
                    CopyOutAxes(i, j, k, tileLength);
                }
                CopyInAxes(tileNum - 1, j, k, totalLength - (tileNum -1)*tileLength);
                Compute(tileNum -1, totalLength - (tileNum -1)*tileLength, lastsum);
                CopyOutAxes(tileNum - 1, j, k, totalLength - (tileNum -1)*tileLength);
            }   
        }
    }


    __aicore__ inline void CopyInAxes(int32_t progress, int j, int k, uint32_t length)
    {
        if(reverse == false)
        {
            LocalTensor<int32_t> xLocal = inQueueX.AllocTensor<int32_t>();
            uint32_t startIndex = k * this->stepSize * this->totalLength +  progress * tileLength * this->stepSize + j;
            for (int i = 0; i < length; i++)
            {
                xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
            }
            inQueueX.EnQue(xLocal);
        }
        else if(reverse == true)
        {
            LocalTensor<int32_t> xLocal = inQueueX.AllocTensor<int32_t>();
            uint32_t startIndex = k * this->stepSize * this->totalLength + (this->totalLength - (progress + 1) * tileLength) * this->stepSize + j;

            if(progress == this->tileNum - 1)   //最后一次tiling
            {   
                startIndex = k * this->stepSize * this->totalLength + j;    // *,*,*,*,0,0
                for (int i = 0; i < length; i++)
                {
                    xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
                }
                for (int i = length; i < this->lasttileLength; i++)     //0 padding
                {
                    xLocal.SetValue(i, (int32_t)0);
                }
            }
            else
            {
                for (int i = 0; i < length; i++)
                {
                    xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
                }
            }
            inQueueX.EnQue(xLocal);
        }
        
    }


    __aicore__ inline void CopyOutAxes(int32_t progress, int j, int k, uint32_t length)
    {
        if(reverse == false)
        {
            uint32_t startIndex = k * this->stepSize * this->totalLength + progress * tileLength * this->stepSize + j;
            LocalTensor<int32_t> yLocal = outQueueY.DeQue<int32_t>();
            for (int i = 0; i < length; i++)
            {
                yGm.SetValue(startIndex + i * stepSize, yLocal.GetValue(i));
            }
            outQueueY.FreeTensor(yLocal);
        }
        else if (reverse == true)
        {
            LocalTensor<int32_t> yLocal = outQueueY.DeQue<int32_t>();
            uint32_t startIndex;
            if(progress == this->tileNum - 1)   //最后一次tiling
            {
                startIndex = k * this->stepSize * this->totalLength + j;
            }
            else
            {
                startIndex = k * this->stepSize * this->totalLength + (this->totalLength - (progress + 1) * tileLength) * this->stepSize + j;
            }
            for(int i = 0; i < length; i++)
            {
                yGm.SetValue(startIndex + i * stepSize, yLocal.GetValue(i));
            }
            outQueueY.FreeTensor(yLocal);
        }    
    }
};


class KernelCumsum_int8 {

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<TPosition::VECCALC> xfBuf, yfBuf;

    GlobalTensor<int8_t> xGm;
    GlobalTensor<int8_t> yGm;
    
    uint32_t minLength;
    uint32_t ALIGN_NUM;

    uint32_t* shape;
    uint32_t dim;

    uint32_t axesNum;

    uint32_t blockLength;
    uint32_t blockSize;
    uint32_t unitCount;
    uint32_t totalLength; //reduce数据个数
    uint32_t stepSize;
    uint32_t totalNum;  //所有数据个数
    
    uint32_t tileNum; //每组reduce数据计算的循环次数
    uint32_t tileLength; //每次tiling的数据长度
    uint32_t lasttileLength; //最后一次tiling的长度/ uint32_t remain_size;

    bool exclusive;
    bool reverse;

public:
    __aicore__  inline KernelCumsum_int8() {}
    __aicore__  inline  void Init(GM_ADDR x, GM_ADDR axes, GM_ADDR y, uint32_t blockLength, uint32_t blockSize, uint32_t ALIGN_NUM, uint32_t* shape, uint32_t dim, bool exclusive, bool reverse)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->blockLength = blockLength;
        this->blockSize = blockSize;
        this->ALIGN_NUM = ALIGN_NUM;
        this->minLength = this->ALIGN_NUM;
        this->shape = shape;
        this->dim = dim;
        this->exclusive = exclusive;
        this->reverse = reverse;

        this->axesNum = (*(__gm__ int32_t *)axes);
        if(axesNum < 0)
        {
            axesNum = dim + axesNum;
        }
        
        this->totalLength = shape[axesNum];
        
        this->stepSize = 1;
        this->unitCount = 1;
        for(uint32_t i = 0; i < axesNum; i++)
        {
            this->unitCount *= shape[i];
        }

        for(uint32_t i = axesNum+1; i < dim; i++)
        {
            this->stepSize *= shape[i];
        }

        this->totalNum = unitCount * totalLength * stepSize;

        uint32_t totalLengthAligned = (totalLength + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

        this->tileNum = totalLengthAligned / this->blockSize;
        // 
        if(totalLengthAligned % blockSize == 0 || tileNum == 0) //reduce dim对齐之后，可以均分到各个tile，或不足一个tile
        {
            if(tileNum == 0)    //一次tiling的数据个数 大于 totalLengthAligned，也就是tile_num = 0时，tileLength=totalLengthAligned
            {
                this->tileNum = 1;
                this->tileLength = totalLengthAligned;
                this->lasttileLength = this->totalLength;
            }
            else //一次tiling的数据个数 小于 totalLengthAligned时，tileLength取最大值tiling_size * ALIGN_NUM
            {
                this->tileLength = this->blockSize;
                this->lasttileLength = this->tileLength;
            }
        }
        else //totalLengthAligned对齐后不能均分到每次tiling
        {
            this->tileNum = this->tileNum + 1;
            this->tileLength = this->blockSize; //tileLength 取最大值(ub)
            this->lasttileLength = totalLengthAligned % this->tileLength;
        }

        // printf("axesNum : %u\n", axesNum);
        // printf("totalLength : %u\n", totalLength);
        // printf("unitCount : %u\n", unitCount);
        // printf("stepSize : %u\n", stepSize);
        // printf("tileNum : %u\n", tileNum);
        // printf("tileLength : %u\n", tileLength);
        // printf("lasttileLength : %u\n", lasttileLength);
        // printf("totalNum : %u\n", totalNum);
        // printf("exclusive : %d\n", exclusive);
        // printf("reverse : %d\n", reverse);

        ASSERT(tileNum != 0 && "tile num can not be zero!");
    
        
        xGm.SetGlobalBuffer((__gm__ int8_t *)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ int8_t *)y, blockLength);

        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(int8_t));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(int8_t));
        pipe.InitBuffer(xfBuf, this->tileLength * sizeof(half));
        pipe.InitBuffer(yfBuf, this->tileLength * sizeof(half));
    }

    __aicore__ inline void Process()
    {
        if (stepSize != 1)
        {
            // 处理轴 != -1情况, 需要SetValue来设置xlocal
            ProcessAxes();
            return;
        }
        //最后一个维度 stepsize = 1
        uint32_t loopCount = this->tileNum;
        for (int32_t k = 0; k < this->unitCount; k++)
        {
            half lastsum = 0.0;
            for(int32_t i = 0; i < loopCount - 1; i++)
            {
                CopyIn(i,  k, tileLength);
                Compute(i, tileLength, lastsum);
                CopyOut(i, k, tileLength);
            }
            uint32_t lastlength = reverse == false ? lasttileLength : totalLength - (tileNum - 1) * tileLength;
            CopyIn(tileNum - 1,  k, lastlength);
            Compute(tileNum - 1, totalLength - (tileNum - 1) * tileLength, lastsum);
            CopyOut(tileNum - 1, k, lastlength);
        }
    }

    __aicore__ inline void CopyIn(int32_t progress, int k, uint32_t length)
    {
        LocalTensor<int8_t> xLocal = inQueueX.AllocTensor<int8_t>();
        if(reverse == false)
        {
            uint32_t startIndex = k  * this->totalLength + progress * tileLength;
            DataCopy(xLocal, xGm[startIndex], length);
            inQueueX.EnQue(xLocal);
        }
        else if(reverse == true)
        {
            uint32_t startIndex = k  * this->totalLength + (this->totalLength -  (progress + 1) * tileLength);
            // startIndex = startIndex > 0 ? startIndex : 0;

            if(progress == this->tileNum - 1)   //最后一次tiling 且无法datacopy, 先copy再覆盖 ，因为lasttileLength 一定小于tileLength
            {
                startIndex = k  * this->totalLength;    //倒序最后一次tiling的start为 “0”  : *,*,*,*,0,0,0
                DataCopy(xLocal, xGm[startIndex], lasttileLength);
                for (int i = length; i < this->lasttileLength; i++)      //0 padding
                {
                    xLocal.SetValue(i, (int8_t)0);
                }
            }
            else
            {
                DataCopy(xLocal, xGm[startIndex], length);
            }

            inQueueX.EnQue(xLocal);
        }
    }


    __aicore__ inline void Compute(int32_t progress, uint32_t length, half& lastsum)
    {   
        if(reverse == false)
        {
            // printf("reverse == false \n");
            LocalTensor<int8_t> xLocal = inQueueX.DeQue<int8_t>();
            LocalTensor<int8_t> yLocal = outQueueY.AllocTensor<int8_t>();

            LocalTensor<half> yLocal_f = yfBuf.Get<half>();
            LocalTensor<half> xLocal_f = xfBuf.Get<half>();

            Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, length);

            Adds(yLocal_f, xLocal_f, (half)0.0, length);
            
            if(exclusive)
            {
                // printf("exclusive == true \n");
                if(length > 1)
                {
                    yLocal_f.SetValue(0, (half)0.0);
                    yLocal_f.SetValue(1, xLocal_f.GetValue(0));
                    for(int i = 2; i < length; i++)
                    {
                        yLocal_f.SetValue(i, (half)((float)yLocal_f.GetValue(i-1) + (float)xLocal_f.GetValue(i-1)));
                    }
                }
                else    // length = 1
                {
                    yLocal_f.SetValue(0, (half)0);
                }
            }
            else
            {   
                // printf("exclusive == false \n");
                for(int i = 1; i < length; i++)
                {
                    yLocal_f.SetValue(i, (half)((float)yLocal_f.GetValue(i-1) + (float)xLocal_f.GetValue(i)));
                }
            }

            Adds(yLocal_f, yLocal_f, lastsum, length);
            lastsum = yLocal.GetValue(length - 1);
            Cast(yLocal, yLocal_f, RoundMode::CAST_NONE, length);

            outQueueY.EnQue<int8_t>(yLocal);
            inQueueX.FreeTensor(xLocal);
        }
        else if (reverse == true)
        {
            // printf("reverse == true \n");
            LocalTensor<int8_t> xLocal = inQueueX.DeQue<int8_t>();
            LocalTensor<int8_t> yLocal = outQueueY.AllocTensor<int8_t>();

            LocalTensor<half> yLocal_f = yfBuf.Get<half>();
            LocalTensor<half> xLocal_f = xfBuf.Get<half>();

            Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, length);

            Adds(yLocal_f, xLocal_f, (half)0.0, length);
            
            if(exclusive)
            {
                // printf("exclusive == true \n");
                if(length > 1)
                    {
                        yLocal_f.SetValue(length - 1, (half)0.0);
                        yLocal_f.SetValue(length - 2, xLocal_f.GetValue(length - 1));
                        for(int i = length - 3; i > -1; i--)
                        {
                            yLocal_f.SetValue(i, (half)((float)xLocal_f.GetValue(i+1) + (float)yLocal_f.GetValue(i+1)));
                        }
                    }
                    else    // length = 1
                    {
                        yLocal_f.SetValue(length - 1, (half)0.0);
                    }
            }
            else
            {   
                // printf("exclusive == false \n");
                for(int i = length - 2; i > -1; i--)
                {
                    yLocal_f.SetValue(i, (half)((float)yLocal_f.GetValue(i+1) + (float)xLocal_f.GetValue(i)));
                }
            }
            Adds(yLocal_f, yLocal_f, lastsum, length);
            lastsum = yLocal_f.GetValue(0);

            Cast(yLocal, yLocal_f, RoundMode::CAST_NONE, length);
            outQueueY.EnQue<int8_t>(yLocal);
            inQueueX.FreeTensor(xLocal);
        }   
    }


    __aicore__ inline void CopyOut(int32_t progress, int k, uint32_t length)
    {
        if(reverse ==false)
        {
            uint32_t startIndex = k * this->totalLength + progress * tileLength;
            LocalTensor<int8_t> yLocal = outQueueY.DeQue<int8_t>();
            DataCopy(yGm[startIndex], yLocal, length);    
            outQueueY.FreeTensor(yLocal);    
        }
        else if(reverse == true)
        {
            LocalTensor<int8_t> yLocal = outQueueY.DeQue<int8_t>();
            if(progress == this->tileNum - 1)   //最后一次tiling
            {
                
                uint32_t startIndex = k  * this->totalLength;
                if(lasttileLength != length)    //存在0padding的数据，只能setvalue拷贝
                    for(int i = 0; i < length; i++)
                    {
                        yGm.SetValue(startIndex + i, yLocal.GetValue(i));
                    }
                else    //不存在0拷贝的数据
                    DataCopy(yGm[startIndex], yLocal, length);
            }
            else
            {
                uint32_t startIndex = k  * this->totalLength + (this->totalLength -  (progress + 1) * tileLength);
                LocalTensor<int8_t> yLocal = outQueueY.DeQue<int8_t>();
                DataCopy(yGm[startIndex], yLocal, length);
            }
            outQueueY.FreeTensor(yLocal);
        }        
    }


    __aicore__ inline void ProcessAxes()
    {
        uint32_t loopCount = this->tileNum;
        for (int32_t k = 0; k < this->unitCount; k++)
        {
            for (int32_t j = 0; j < this->stepSize; j++)
            {
                half lastsum = 0.0;
                for(int32_t i = 0; i < loopCount - 1; i++)
                {
                    CopyInAxes(i, j, k, tileLength);
                    Compute(i, tileLength, lastsum);
                    CopyOutAxes(i, j, k, tileLength);
                }
                CopyInAxes(tileNum - 1, j, k, totalLength - (tileNum -1)*tileLength);
                Compute(tileNum -1, totalLength - (tileNum -1)*tileLength, lastsum);
                CopyOutAxes(tileNum - 1, j, k, totalLength - (tileNum -1)*tileLength);
            }   
        }
    }


    __aicore__ inline void CopyInAxes(int32_t progress, int j, int k, uint32_t length)
    {
        if(reverse == false)
        {
            LocalTensor<int8_t> xLocal = inQueueX.AllocTensor<int8_t>();
            uint32_t startIndex = k * this->stepSize * this->totalLength +  progress * tileLength * this->stepSize + j;
            for (int i = 0; i < length; i++)
            {
                xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
            }
            inQueueX.EnQue(xLocal);
        }
        else if(reverse == true)
        {
            LocalTensor<int8_t> xLocal = inQueueX.AllocTensor<int8_t>();
            uint32_t startIndex = k * this->stepSize * this->totalLength + (this->totalLength - (progress + 1) * tileLength) * this->stepSize + j;

            if(progress == this->tileNum - 1)   //最后一次tiling
            {   
                startIndex = k * this->stepSize * this->totalLength + j;    // *,*,*,*,0,0
                for (int i = 0; i < length; i++)
                {
                    xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
                }
                for (int i = length; i < this->lasttileLength; i++)     //0 padding
                {
                    xLocal.SetValue(i, (int8_t)0);
                }
            }
            else
            {
                for (int i = 0; i < length; i++)
                {
                    xLocal.SetValue(i, xGm.GetValue(startIndex + i * stepSize));
                }
            }
            inQueueX.EnQue(xLocal);
        }        
    }

    __aicore__ inline void CopyOutAxes(int32_t progress, int j, int k, uint32_t length)
    {
        if(reverse == false)
        {
            uint32_t startIndex = k * this->stepSize * this->totalLength + progress * tileLength * this->stepSize + j;
            LocalTensor<int8_t> yLocal = outQueueY.DeQue<int8_t>();
            for (int i = 0; i < length; i++)
            {
                yGm.SetValue(startIndex + i * stepSize, yLocal.GetValue(i));
            }
            outQueueY.FreeTensor(yLocal);
        }
        else if (reverse == true)
        {
            LocalTensor<int8_t> yLocal = outQueueY.DeQue<int8_t>();
            uint32_t startIndex;
            if(progress == this->tileNum - 1)   //最后一次tiling
            {
                startIndex = k * this->stepSize * this->totalLength + j;
            }
            else
            {
                startIndex = k * this->stepSize * this->totalLength + (this->totalLength - (progress + 1) * tileLength) * this->stepSize + j;
            }
            for(int i = 0; i < length; i++)
            {
                yGm.SetValue(startIndex + i * stepSize, yLocal.GetValue(i));
            }
            outQueueY.FreeTensor(yLocal);
        }    
    }
};



extern "C" __global__ __aicore__ void cumsum(GM_ADDR x, GM_ADDR axis, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    if(TILING_KEY_IS(1))
    {
        KernelCumsum_fp32 op;
        op.Init(x, axis, y, tiling_data.blockLength, tiling_data.blockSize, tiling_data.ALIGN_NUM, tiling_data.shape, tiling_data.dimNum, tiling_data.exclusive, tiling_data.reverse);
        op.Process();
    }
    else if (TILING_KEY_IS(2))
    {
        KernelCumsum_fp16 op;
        op.Init(x, axis, y, tiling_data.blockLength, tiling_data.blockSize, tiling_data.ALIGN_NUM, tiling_data.shape, tiling_data.dimNum, tiling_data.exclusive, tiling_data.reverse);
        op.Process();
    }
    else if (TILING_KEY_IS(3))
    {
        KernelCumsum_int32 op;
        op.Init(x, axis, y, tiling_data.blockLength, tiling_data.blockSize, tiling_data.ALIGN_NUM, tiling_data.shape, tiling_data.dimNum, tiling_data.exclusive, tiling_data.reverse);
        op.Process();
    }
    else if (TILING_KEY_IS(4))
    {
        KernelCumsum_int8 op;
        op.Init(x, axis, y, tiling_data.blockLength, tiling_data.blockSize, tiling_data.ALIGN_NUM, tiling_data.shape, tiling_data.dimNum, tiling_data.exclusive, tiling_data.reverse);
        op.Process();
    }
}